import subprocess
import sys
import os
import tempfile
from pathlib import Path

def create_conda_env():
    """创建 Conda 虚拟环境并安装指定的 pip 包，支持环境检查和更新"""
    # 直接在代码中指定参数
    ENV_NAME = "myenv"               # 环境名称
    PYTHON_VERSION = "3.10.8"        # Python 版本
    REQUIREMENTS_FILE = "requirements.txt"  # 依赖文件
    
    # 检查 Conda 是否已安装
    try:
        subprocess.run(["conda", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("错误: 未找到 Conda。请确保已安装 Anaconda 或 Miniconda。")
        sys.exit(1)

    # 检查环境是否已存在
    env_exists = subprocess.run(
        f"conda env list | grep -q '^{ENV_NAME}\\s'", 
        shell=True, 
        stdout=subprocess.PIPE, 
        stderr=subprocess.PIPE
    ).returncode == 0
    
    if env_exists:
        print(f"检测到 Conda 环境 '{ENV_NAME}' 已存在")
        
        # 检查 requirements 文件是否存在
        if not os.path.exists(REQUIREMENTS_FILE):
            print(f"警告: 未找到 requirements 文件 '{REQUIREMENTS_FILE}'，跳过包检查。")
            return
            
        # 检查环境中的包与 requirements 文件是否一致
        missing_packages = check_missing_packages(ENV_NAME, REQUIREMENTS_FILE)
        
        if not missing_packages:
            print(f"环境 '{ENV_NAME}' 已安装 requirements 文件中的所有包")
            return
            
        print(f"\n环境 '{ENV_NAME}' 缺少以下包:")
        for pkg in missing_packages:
            print(f"  - {pkg}")
            
        # 询问用户是否更新环境
        answer = 'y'#input("\n是否更新环境以安装缺失的包? [y/N]: ").strip().lower()
        if answer != 'y':
            print("已取消更新")
            return
            
        # 更新环境
        update_existing_env(ENV_NAME, REQUIREMENTS_FILE)
        return
    
    # 创建新环境
    print(f"正在创建 Conda 环境 '{ENV_NAME}'，Python 版本: {PYTHON_VERSION}")
    create_cmd = [
        "conda", "create", "-n", ENV_NAME,
        f"python={PYTHON_VERSION}", "-y"
    ]
    
    # 实时显示 conda 创建环境的输出
    process = subprocess.Popen(
        create_cmd, 
        stdout=subprocess.PIPE, 
        stderr=subprocess.STDOUT,
        universal_newlines=True
    )
    
    for line in iter(process.stdout.readline, ''):
        print(line, end='')
    
    returncode = process.wait()
    if returncode != 0:
        print(f"\n创建环境失败，返回代码: {returncode}")
        sys.exit(1)
    
    print(f"\n已成功创建 Conda 环境 '{ENV_NAME}'")
    
    # 检查 requirements 文件是否存在
    if not os.path.exists(REQUIREMENTS_FILE):
        print(f"警告: 未找到 requirements 文件 '{REQUIREMENTS_FILE}'，跳过 pip 包安装。")
        return
    
    # 在新环境中安装 pip 包
    install_packages(ENV_NAME, REQUIREMENTS_FILE)

def check_missing_packages(env_name, requirements_file):
    """检查环境中缺少的包"""
    print(f"正在检查环境 '{env_name}' 中缺少的包...")
    
    # 获取环境中已安装的 pip 包
    list_cmd = f"""
    source $(conda info --base)/etc/profile.d/conda.sh &&
    conda activate {env_name} &&
    pip list --format=freeze
    """
    
    result = subprocess.run(
        list_cmd, 
        shell=True, 
        executable="/bin/bash", 
        capture_output=True, 
        text=True
    )
    
    if result.returncode != 0:
        print(f"获取已安装包列表失败: {result.stderr.strip()}")
        sys.exit(1)
    
    installed_packages = set(result.stdout.strip().split('\n'))
    
    # 读取 requirements 文件
    with open(requirements_file, 'r') as f:
        required_packages = set(line.strip() for line in f if line.strip() and not line.startswith('#'))
    
    # 找出缺失的包
    missing_packages = []
    
    for req in required_packages:
        # 处理包名和版本约束
        if '==' in req:
            pkg_name = req.split('==')[0]
            # 检查包名和版本是否完全匹配
            if not any(pkg.startswith(req) for pkg in installed_packages):
                missing_packages.append(req)
        else:
            # 只检查包名是否存在
            if not any(pkg.split('==')[0] == req for pkg in installed_packages):
                missing_packages.append(req)
    
    return missing_packages

def install_packages(env_name, requirements_file):
    """在指定环境中安装 requirements 文件中的包"""
    print(f"\n正在从 '{requirements_file}' 安装 pip 包到环境 '{env_name}'...")
    
    install_cmd = f"""
    source $(conda info --base)/etc/profile.d/conda.sh &&
    conda activate {env_name} &&
    pip install -r {requirements_file}
    """
    
    # 实时显示 pip 安装的输出
    process = subprocess.Popen(
        install_cmd, 
        shell=True, 
        executable="/bin/bash", 
        stdout=subprocess.PIPE, 
        stderr=subprocess.STDOUT,
        universal_newlines=True
    )
    
    for line in iter(process.stdout.readline, ''):
        print(line, end='')
    
    returncode = process.wait()
    if returncode != 0:
        print(f"\n安装 pip 包失败，返回代码: {returncode}")
        sys.exit(1)
    
    print(f"\n已成功安装所有 pip 包到环境 '{env_name}'")
    print("\n环境创建完成! 使用以下命令激活环境:")
    print(f"  conda activate {env_name}")

def update_existing_env(env_name, requirements_file):
    """更新现有环境，安装缺失的包"""
    print(f"\n正在更新环境 '{env_name}'...")
    
    install_cmd = f"""
    source $(conda info --base)/etc/profile.d/conda.sh &&
    conda activate {env_name} &&
    pip install -r {requirements_file}
    """
    
    # 实时显示 pip 安装的输出
    process = subprocess.Popen(
        install_cmd, 
        shell=True, 
        executable="/bin/bash", 
        stdout=subprocess.PIPE, 
        stderr=subprocess.STDOUT,
        universal_newlines=True
    )
    
    for line in iter(process.stdout.readline, ''):
        print(line, end='')
    
    returncode = process.wait()
    if returncode != 0:
        print(f"\n更新环境失败，返回代码: {returncode}")
        sys.exit(1)
    
    print(f"\n已成功更新环境 '{env_name}'")
    print("\n使用以下命令激活环境:")
    print(f"  conda activate {env_name}")

if __name__ == "__main__":
    create_conda_env()
